{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9bc8f086-5dda-4aee-9dcb-14e4d9720f0d",
   "metadata": {},
   "source": [
    "# Fast Dequantizatoin\n",
    "\n",
    "How to enbale fast dequantization (INT4/2/1 -> FP16/INT8) or (FP8 -> FP16)?\n",
    "\n",
    "![image.png](./img/FastDequantization.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "32b3eb3a-a0a2-4c63-9e24-a248c1fe3579",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-10-24 14:35:48 [BitBLAS:INFO]: Auto detected target: nvidia/geforce-rtx-4090\n",
      "2024-10-24 14:35:49 [BitBLAS:DEBUG]: Cannot find the appropriate index map for tensorcore\n",
      "2024-10-24 14:36:16 [BitBLAS:DEBUG]: Cannot find the appropriate index map for tensorcore\n",
      "2024-10-24 14:36:16 [BitBLAS:DEBUG]: Apply config {'block': [2], 'thread': [2], 'rstep': [1024], 'reduce_thread': [64], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:36:16 [BitBLAS:DEBUG]: Apply config {'block': [1], 'thread': [1], 'rstep': [1024], 'reduce_thread': [128], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:36:16 [BitBLAS:DEBUG]: Apply config {'block': [4], 'thread': [4], 'rstep': [1024], 'reduce_thread': [32], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:36:16 [BitBLAS:DEBUG]: Apply config {'block': [8], 'thread': [8], 'rstep': [1024], 'reduce_thread': [16], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:36:16 [BitBLAS:DEBUG]: Apply config {'block': [16], 'thread': [16], 'rstep': [512], 'reduce_thread': [8], 'vectorize': {'A': 4, 'B_decode': 8}}\n",
      "2024-10-24 14:36:16 [BitBLAS:DEBUG]: Apply config {'block': [32], 'thread': [32], 'rstep': [256], 'reduce_thread': [4], 'vectorize': {'A': 2, 'B_decode': 8}}\n",
      "2024-10-24 14:36:16 [BitBLAS:DEBUG]: Apply config {'block': [64], 'thread': [64], 'rstep': [128], 'reduce_thread': [2], 'vectorize': {'B_decode': 8}}\n",
      "2024-10-24 14:36:16 [BitBLAS:DEBUG]: Apply config {'block': [128], 'thread': [128], 'rstep': [128], 'vectorize': {'B_decode': 8}}\n",
      "2024-10-24 14:36:17 [BitBLAS:DEBUG]: Warning: block config [128] is not valid for matmul, skip.\n",
      "2024-10-24 14:36:17 [BitBLAS:DEBUG]: Warning: block config [128] is not valid for matmul, skip.\n",
      "2024-10-24 14:36:42 [BitBLAS:DEBUG]: LocalBuilder: An exception occurred Traceback (most recent call last):\n",
      "  File \"/root/BitBLAS/3rdparty/tvm/python/tvm/exec/popen_worker.py\", line 87, in main\n",
      "    result = fn(*args, **kwargs)\n",
      "  File \"/root/BitBLAS/bitblas/base/utils.py\", line 257, in _build\n",
      "    rt_mod = tvm.build(mod, ta\t...\tm::codegen::CodeGenCUDA::VisitExpr_(tvm::tir::RampNode const*, std::ostream&)\n",
      "  File \"/root/BitBLAS/3rdparty/tvm/src/target/source/codegen_cuda.cc\", line 1226\n",
      "ValueError: Check failed: lanes <= 4 (8 vs. 4) : Ramp of more than 4 lanes is not allowed.\n",
      "\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Evaluation with config {'block': [2], 'thread': [2], 'rstep': [1024], 'reduce_thread': [64], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Time cost of this config: 0.005 ms\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Evaluation with config {'block': [1], 'thread': [1], 'rstep': [1024], 'reduce_thread': [128], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Time cost of this config: 0.004 ms\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Evaluation with config {'block': [4], 'thread': [4], 'rstep': [1024], 'reduce_thread': [32], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Time cost of this config: 0.004 ms\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Evaluation with config {'block': [8], 'thread': [8], 'rstep': [1024], 'reduce_thread': [16], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Time cost of this config: 0.004 ms\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Evaluation with config {'block': [16], 'thread': [16], 'rstep': [512], 'reduce_thread': [8], 'vectorize': {'A': 4, 'B_decode': 8}}\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Time cost of this config: 0.004 ms\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Evaluation with config {'block': [32], 'thread': [32], 'rstep': [256], 'reduce_thread': [4], 'vectorize': {'A': 2, 'B_decode': 8}}\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Time cost of this config: 0.005 ms\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Evaluation with config {'block': [64], 'thread': [64], 'rstep': [128], 'reduce_thread': [2], 'vectorize': {'B_decode': 8}}\n",
      "2024-10-24 14:36:42 [BitBLAS:INFO]: Time cost of this config: 0.008 ms\n",
      "Ref output: tensor([[1586., 1519., 1561.,  ..., 1566., 1536., 1561.]], device='cuda:0',\n",
      "       dtype=torch.float16)\n",
      "BitBLAS output: tensor([[1586., 1518., 1562.,  ..., 1567., 1535., 1562.]], device='cuda:0',\n",
      "       dtype=torch.float16)\n"
     ]
    }
   ],
   "source": [
    "import bitblas\n",
    "import torch\n",
    "\n",
    "# enabling debug output\n",
    "\n",
    "bitblas.set_log_level(\"Debug\")\n",
    "matmul_config = bitblas.MatmulConfig(\n",
    "    M=1,  # M dimension\n",
    "    N=1024,  # N dimension\n",
    "    K=1024,  # K dimension\n",
    "    A_dtype=\"float16\",  # activation A dtype\n",
    "    W_dtype=\"int4\",  # weight W dtype\n",
    "    accum_dtype=\"float16\",  # accumulation dtype\n",
    "    out_dtype=\"float16\",  # output dtype\n",
    "    layout=\"nt\",  # matrix layout, \"nt\" indicates the layout of A is non-transpose and the layout of W is transpose\n",
    "    with_bias=False,  # bias\n",
    "    # configs for weight only quantization\n",
    "    group_size=None,  # setting for grouped quantization\n",
    "    with_scaling=False,  # setting for scaling factor\n",
    "    with_zeros=False,  # setting for zeros\n",
    "    zeros_mode=None,  # setting for how to calculating zeros\n",
    ")\n",
    "\n",
    "matmul = bitblas.Matmul(config=matmul_config)\n",
    "\n",
    "# Create input matrices\n",
    "input_tensor = torch.rand((1, 1024), dtype=torch.float16).cuda()\n",
    "weight_tensor = torch.randint(0, 7, (1024, 1024), dtype=torch.int8).cuda()\n",
    "\n",
    "# Transform weight tensor to int4 data type\n",
    "weight_tensor_int4 = matmul.transform_weight(weight_tensor)\n",
    "\n",
    "# Perform mixed-precision matrix multiplication\n",
    "output_tensor = matmul(input_tensor, weight_tensor_int4)\n",
    "\n",
    "# Reference result using PyTorch matmul for comparison\n",
    "ref_result = torch.matmul(input_tensor, weight_tensor.t().to(torch.float16))\n",
    "# Assert that the results are close within a specified tolerance, note that the int4 randint value is a little bigger than the float16 value, so we set the atol to 1.0\n",
    "print(\"Ref output:\", ref_result)\n",
    "print(\"BitBLAS output:\", output_tensor)\n",
    "torch.testing.assert_close(output_tensor, ref_result, rtol=1e-2, atol=1e-0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "839bd993-0120-4061-9799-88b55f905ce1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import bitblas\n",
    "import torch\n",
    "\n",
    "# enabling debug output\n",
    "bitblas.set_log_level(\"Debug\")\n",
    "matmul_config = bitblas.MatmulConfig(\n",
    "    M=1,  # M dimension\n",
    "    N=16384,  # N dimension\n",
    "    K=16384,  # K dimension\n",
    "    A_dtype=\"float16\",  # activation A dtype\n",
    "    W_dtype=\"float16\",  # weight W dtype\n",
    ")\n",
    "\n",
    "matmul = bitblas.Matmul(config=matmul_config, enable_tuning=False)\n",
    "\n",
    "matmul.hardware_aware_finetune(topk=20, parallel_build=True)\n",
    "\n",
    "latency = matmul.profile_latency()\n",
    "print(f\"{latency=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "832054b0-35a6-4422-b008-c6759e3a9162",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-10-24 14:45:08 [BitBLAS:INFO]: Auto detected target: nvidia/geforce-rtx-4090\n",
      "2024-10-24 14:45:08 [BitBLAS:DEBUG]: Cannot find the appropriate index map for tensorcore\n",
      "2024-10-24 14:45:10 [BitBLAS:DEBUG]: Cannot find the appropriate index map for tensorcore\n",
      "2024-10-24 14:45:11 [BitBLAS:DEBUG]: Apply config {'block': [32], 'thread': [32], 'rstep': [256], 'reduce_thread': [4], 'vectorize': {'A': 2, 'B_decode': 8}}\n",
      "2024-10-24 14:45:11 [BitBLAS:DEBUG]: Apply config {'block': [16], 'thread': [16], 'rstep': [512], 'reduce_thread': [8], 'vectorize': {'A': 4, 'B_decode': 8}}\n",
      "2024-10-24 14:45:11 [BitBLAS:DEBUG]: Apply config {'block': [8], 'thread': [8], 'rstep': [1024], 'reduce_thread': [16], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:45:11 [BitBLAS:DEBUG]: Apply config {'block': [4], 'thread': [4], 'rstep': [2048], 'reduce_thread': [32], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:45:11 [BitBLAS:DEBUG]: Apply config {'block': [2], 'thread': [2], 'rstep': [4096], 'reduce_thread': [64], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:45:11 [BitBLAS:DEBUG]: Apply config {'block': [1], 'thread': [1], 'rstep': [4096], 'reduce_thread': [128], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:45:11 [BitBLAS:DEBUG]: Apply config {'block': [64], 'thread': [64], 'rstep': [128], 'reduce_thread': [2], 'vectorize': {'B_decode': 8}}\n",
      "2024-10-24 14:45:11 [BitBLAS:DEBUG]: Apply config {'block': [128], 'thread': [128], 'rstep': [128], 'vectorize': {'B_decode': 8}}\n",
      "2024-10-24 14:45:11 [BitBLAS:DEBUG]: Warning: block config [128] is not valid for matmul, skip.\n",
      "2024-10-24 14:45:11 [BitBLAS:DEBUG]: Warning: block config [128] is not valid for matmul, skip.\n"
     ]
    }
   ],
   "source": [
    "matmul_config = bitblas.MatmulConfig(\n",
    "    M=1,  # M dimension\n",
    "    N=16384,  # N dimension\n",
    "    K=16384,  # K dimension\n",
    "    A_dtype=\"float16\",  # activation A dtype\n",
    "    W_dtype=\"int4\",  # weight W dtype\n",
    "    fast_decoding=False, # Disable Fast Decoding\n",
    ")\n",
    "\n",
    "matmul = bitblas.Matmul(config=matmul_config, enable_tuning=False)\n",
    "\n",
    "matmul.hardware_aware_finetune(topk=20, parallel_build=True)\n",
    "\n",
    "print(matmul.get_source())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5da4e137-5d02-46e7-9cb0-185d8d0137f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "latency=0.1433375\n"
     ]
    }
   ],
   "source": [
    "latency = matmul.profile_latency()\n",
    "print(f\"{latency=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b7b4a317-9550-42df-8c3e-be88069eabf8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-10-24 14:44:22 [BitBLAS:INFO]: Auto detected target: nvidia/geforce-rtx-4090\n",
      "2024-10-24 14:44:22 [BitBLAS:DEBUG]: Cannot find the appropriate index map for tensorcore\n",
      "2024-10-24 14:44:25 [BitBLAS:DEBUG]: Cannot find the appropriate index map for tensorcore\n",
      "2024-10-24 14:44:25 [BitBLAS:DEBUG]: Apply config {'block': [32], 'thread': [32], 'rstep': [256], 'reduce_thread': [4], 'vectorize': {'A': 2, 'B_decode': 8}}\n",
      "2024-10-24 14:44:25 [BitBLAS:DEBUG]: Apply config {'block': [16], 'thread': [16], 'rstep': [512], 'reduce_thread': [8], 'vectorize': {'A': 4, 'B_decode': 8}}\n",
      "2024-10-24 14:44:25 [BitBLAS:DEBUG]: Apply config {'block': [8], 'thread': [8], 'rstep': [1024], 'reduce_thread': [16], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:44:25 [BitBLAS:DEBUG]: Apply config {'block': [4], 'thread': [4], 'rstep': [2048], 'reduce_thread': [32], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:44:25 [BitBLAS:DEBUG]: Apply config {'block': [2], 'thread': [2], 'rstep': [4096], 'reduce_thread': [64], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:44:25 [BitBLAS:DEBUG]: Apply config {'block': [1], 'thread': [1], 'rstep': [4096], 'reduce_thread': [128], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:44:25 [BitBLAS:DEBUG]: Apply config {'block': [64], 'thread': [64], 'rstep': [128], 'reduce_thread': [2], 'vectorize': {'B_decode': 8}}\n",
      "2024-10-24 14:44:25 [BitBLAS:DEBUG]: Apply config {'block': [128], 'thread': [128], 'rstep': [128], 'vectorize': {'B_decode': 8}}\n",
      "2024-10-24 14:44:25 [BitBLAS:DEBUG]: Warning: block config [128] is not valid for matmul, skip.\n",
      "2024-10-24 14:44:25 [BitBLAS:DEBUG]: Warning: block config [128] is not valid for matmul, skip.\n",
      "2024-10-24 14:44:51 [BitBLAS:DEBUG]: LocalBuilder: An exception occurred Traceback (most recent call last):\n",
      "  File \"/root/BitBLAS/3rdparty/tvm/python/tvm/exec/popen_worker.py\", line 87, in main\n",
      "    result = fn(*args, **kwargs)\n",
      "  File \"/root/BitBLAS/bitblas/base/utils.py\", line 257, in _build\n",
      "    rt_mod = tvm.build(mod, ta\t...\tm::codegen::CodeGenCUDA::VisitExpr_(tvm::tir::RampNode const*, std::ostream&)\n",
      "  File \"/root/BitBLAS/3rdparty/tvm/src/target/source/codegen_cuda.cc\", line 1226\n",
      "ValueError: Check failed: lanes <= 4 (8 vs. 4) : Ramp of more than 4 lanes is not allowed.\n",
      "\n",
      "2024-10-24 14:44:52 [BitBLAS:INFO]: Evaluation with config {'block': [32], 'thread': [32], 'rstep': [256], 'reduce_thread': [4], 'vectorize': {'A': 2, 'B_decode': 8}}\n",
      "2024-10-24 14:44:52 [BitBLAS:INFO]: Time cost of this config: 0.151 ms\n",
      "2024-10-24 14:44:53 [BitBLAS:INFO]: Evaluation with config {'block': [16], 'thread': [16], 'rstep': [512], 'reduce_thread': [8], 'vectorize': {'A': 4, 'B_decode': 8}}\n",
      "2024-10-24 14:44:53 [BitBLAS:INFO]: Time cost of this config: 0.144 ms\n",
      "2024-10-24 14:44:54 [BitBLAS:INFO]: Evaluation with config {'block': [8], 'thread': [8], 'rstep': [1024], 'reduce_thread': [16], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:44:54 [BitBLAS:INFO]: Time cost of this config: 0.161 ms\n",
      "2024-10-24 14:44:55 [BitBLAS:INFO]: Evaluation with config {'block': [4], 'thread': [4], 'rstep': [2048], 'reduce_thread': [32], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:44:55 [BitBLAS:INFO]: Time cost of this config: 0.145 ms\n",
      "2024-10-24 14:44:57 [BitBLAS:INFO]: Evaluation with config {'block': [2], 'thread': [2], 'rstep': [4096], 'reduce_thread': [64], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:44:57 [BitBLAS:INFO]: Time cost of this config: 0.628 ms\n",
      "2024-10-24 14:44:58 [BitBLAS:INFO]: Evaluation with config {'block': [1], 'thread': [1], 'rstep': [4096], 'reduce_thread': [128], 'vectorize': {'A': 8, 'B_decode': 8}}\n",
      "2024-10-24 14:44:58 [BitBLAS:INFO]: Time cost of this config: 0.305 ms\n",
      "2024-10-24 14:44:59 [BitBLAS:INFO]: Evaluation with config {'block': [128], 'thread': [128], 'rstep': [128], 'vectorize': {'B_decode': 8}}\n",
      "2024-10-24 14:44:59 [BitBLAS:INFO]: Time cost of this config: 0.144 ms\n",
      "#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 530)\n",
      "\n",
      "typedef unsigned short uint16_t;\n",
      "typedef unsigned char uint8_t;\n",
      "typedef signed char int8_t;\n",
      "typedef int int32_t;\n",
      "typedef unsigned long long uint64_t;\n",
      "typedef unsigned int uint32_t;\n",
      "\n",
      "#define TVM_FORCE_INLINE inline __attribute__((always_inline))\n",
      "#define TVM_XINLINE TVM_FORCE_INLINE __device__ __host__\n",
      "#define TVM_ALIGNED(x) __attribute__ ((aligned(x)))\n",
      "#define TVM_HALF_OPERATOR(RTYPE, OP)                              \\\n",
      "  TVM_XINLINE RTYPE operator OP (half a, half b) {                \\\n",
      "    return RTYPE(float(a) OP float(b));                           \\\n",
      "  }                                                               \\\n",
      "  template<typename T>                                            \\\n",
      "  TVM_XINLINE RTYPE operator OP (half a, T b) {                   \\\n",
      "    return RTYPE(float(a) OP float(b));                           \\\n",
      "  }                                                               \\\n",
      "  template<typename T>                                            \\\n",
      "  TVM_XINLINE RTYPE operator OP (T a, half b) {                   \\\n",
      "    return RTYPE(float(a) OP float(b));                           \\\n",
      "  }\n",
      "\n",
      "#define TVM_HALF_ASSIGNOP(AOP, OP)                                \\\n",
      "  template<typename T>                                            \\\n",
      "  TVM_XINLINE half operator AOP (const T& a) {                    \\\n",
      "    return *this = half(float(*this) OP float(a));                \\\n",
      "  }                                                               \\\n",
      "  template<typename T>                                            \\\n",
      "  TVM_XINLINE half operator AOP (const volatile T& a) volatile {  \\\n",
      "    return *this = half(float(*this) OP float(a));                \\\n",
      "  }\n",
      "\n",
      "class TVM_ALIGNED(2) half {\n",
      " public:\n",
      "  uint16_t half_;\n",
      "\n",
      "  static TVM_XINLINE half Binary(uint16_t value) {\n",
      "    half res;\n",
      "    res.half_ = value;\n",
      "    return res;\n",
      "  }\n",
      "\n",
      "  TVM_XINLINE half() {}\n",
      "\n",
      "  TVM_XINLINE half(const float& value) { constructor(value); }\n",
      "  TVM_XINLINE explicit half(const double& value) { constructor(value); }\n",
      "  TVM_XINLINE explicit half(const int8_t& value) { constructor(value); }\n",
      "  TVM_XINLINE explicit half(const uint8_t& value) { constructor(value); }\n",
      "  TVM_XINLINE explicit half(const int32_t& value) { constructor(value); }\n",
      "  TVM_XINLINE explicit half(const uint32_t& value) { constructor(value); }\n",
      "  TVM_XINLINE explicit half(const long long& value) { constructor(value); }\n",
      "  TVM_XINLINE explicit half(const uint64_t& value) { constructor(value); }\n",
      "\n",
      "  TVM_XINLINE operator float() const {                          \\\n",
      "    return float(half2float(half_));                            \\\n",
      "  }                                                             \\\n",
      "  TVM_XINLINE operator float() const volatile {                 \\\n",
      "    return float(half2float(half_));                            \\\n",
      "  }\n",
      "\n",
      "\n",
      "  TVM_HALF_ASSIGNOP(+=, +)\n",
      "  TVM_HALF_ASSIGNOP(-=, -)\n",
      "  TVM_HALF_ASSIGNOP(*=, *)\n",
      "  TVM_HALF_ASSIGNOP(/=, /)\n",
      "\n",
      "  TVM_XINLINE half operator+() {\n",
      "    return *this;\n",
      "  }\n",
      "\n",
      "  TVM_XINLINE half operator-() {\n",
      "    return half(-float(*this));\n",
      "  }\n",
      "\n",
      "  TVM_XINLINE half operator=(const half& a) {\n",
      "    half_ = a.half_;\n",
      "    return a;\n",
      "  }\n",
      "\n",
      "  template<typename T>\n",
      "  TVM_XINLINE half operator=(const T& a) {\n",
      "    return *this = half(a);\n",
      "  }\n",
      "\n",
      "  TVM_XINLINE half operator=(const half& a) volatile {\n",
      "    half_ = a.half_;\n",
      "    return a;\n",
      "  }\n",
      "\n",
      "  template<typename T>\n",
      "  TVM_XINLINE half operator=(const T& a) volatile {\n",
      "    return *this = half(a);\n",
      "  }\n",
      "\n",
      " private:\n",
      "  union Bits {\n",
      "    float f;\n",
      "    int32_t si;\n",
      "    uint32_t ui;\n",
      "  };\n",
      "\n",
      "  static int const fp16FractionBits = 10;\n",
      "  static int const fp32FractionBits = 23;\n",
      "  static int32_t const fp32FractionMask = ~(~0u << fp32FractionBits);   // == 0x7fffff\n",
      "  static int32_t const fp32HiddenBit = 1 << fp32FractionBits;   // == 0x800000\n",
      "  static int const shift = fp32FractionBits - fp16FractionBits;   // == 13\n",
      "  static int const shiftSign = 16;\n",
      "  static int32_t const expAdjust = 127 - 15;   // exp32-127 = exp16-15, so exp16 = exp32 - (127-15)\n",
      "\n",
      "  static int32_t const infN = 0x7F800000;   // flt32 infinity\n",
      "  static int32_t const maxN = 0x477FFFFF;   // max flt32 that's a flt16 normal after >> by shift\n",
      "  static int32_t const minN = 0x38800000;   // min flt16 normal as a flt32\n",
      "  static int32_t const maxZ = 0x33000000;   // max fp32 number that's still rounded to zero in fp16\n",
      "  static int32_t const signN = 0x80000000;  // flt32 sign bit\n",
      "\n",
      "  static int32_t const infC = infN >> shift;\n",
      "  static int32_t const nanN = (infC + 1) << shift;   // minimum flt16 nan as a flt32\n",
      "  static int32_t const maxC = maxN >> shift;\n",
      "  static int32_t const minC = minN >> shift;\n",
      "  static int32_t const signC = signN >> shiftSign;  // flt16 sign bit\n",
      "\n",
      "  static int32_t const mulN = 0x52000000;  // (1 << 23) / minN\n",
      "  static int32_t const mulC = 0x33800000;  // minN / (1 << (23 - shift))\n",
      "\n",
      "  static int32_t const subC = 0x003FF;  // max flt32 subnormal down shifted\n",
      "  static int32_t const norC = 0x00400;  // min flt32 normal down shifted\n",
      "\n",
      "  static int32_t const maxD = infC - maxC - 1;\n",
      "  static int32_t const minD = minC - subC - 1;\n",
      "\n",
      "  TVM_XINLINE uint16_t float2half(const float& value) const {\n",
      "    Bits v;\n",
      "    v.f = value;\n",
      "    uint32_t sign = v.si & signN;    // grab sign bit\n",
      "    v.si ^= sign;                    // clear sign bit from v\n",
      "    sign >>= shiftSign;              // logical shift sign to fp16 position\n",
      "\n",
      "    if (v.si <= maxZ) {\n",
      "      // Handle eventual zeros here to ensure\n",
      "      // vshift will not exceed 32 below.\n",
      "      v.ui = 0;\n",
      "    } else if (v.si < minN) {\n",
      "      // Handle denorms\n",
      "      uint32_t exp32 = v.ui >> fp32FractionBits;\n",
      "      int32_t exp16 = exp32 - expAdjust;\n",
      "      // If exp16 == 0 (just into the denorm range), then significant should be shifted right 1.\n",
      "      // Smaller (so negative) exp16 values should result in greater right shifts.\n",
      "      uint32_t vshift = 1 - exp16;\n",
      "      uint32_t significand = fp32HiddenBit | (v.ui & fp32FractionMask);\n",
      "      v.ui = significand >> vshift;\n",
      "      v.ui += (v.ui & 0x3fff) != 0x1000 || (significand & 0x7ff) ? 0x1000 : 0;\n",
      "    } else if (v.si <= maxN) {\n",
      "      // Handle norms\n",
      "      v.ui += (v.ui & 0x3fff) != 0x1000 ? 0x1000 : 0;\n",
      "      v.ui -= expAdjust << fp32FractionBits;\n",
      "    } else if (v.si <= infN) {\n",
      "      v.si = infN;\n",
      "    } else if (v.si < nanN) {\n",
      "      v.si = nanN;\n",
      "    }\n",
      "\n",
      "    v.ui >>= shift;\n",
      "    return sign | (v.ui & 0x7fff);\n",
      "  }\n",
      "\n",
      "  // Same as above routine, except for addition of volatile keyword\n",
      "  TVM_XINLINE uint16_t float2half(\n",
      "    const volatile float& value) const volatile {\n",
      "    Bits v;\n",
      "    v.f = value;\n",
      "    uint32_t sign = v.si & signN;    // grab sign bit\n",
      "    v.si ^= sign;                    // clear sign bit from v\n",
      "    sign >>= shiftSign;              // logical shift sign to fp16 position\n",
      "\n",
      "    if (v.si <= maxZ) {\n",
      "      // Handle eventual zeros here to ensure\n",
      "      // vshift will not exceed 32 below.\n",
      "      v.ui = 0;\n",
      "    } else if (v.si < minN) {\n",
      "      // Handle denorms\n",
      "      uint32_t exp32 = v.ui >> fp32FractionBits;\n",
      "      int32_t exp16 = exp32 - expAdjust;\n",
      "      // If exp16 == 0 (just into the denorm range), then significant should be shifted right 1.\n",
      "      // Smaller (so negative) exp16 values should result in greater right shifts.\n",
      "      uint32_t vshift = 1 - exp16;\n",
      "      uint32_t significand = fp32HiddenBit | (v.ui & fp32FractionMask);\n",
      "      v.ui = significand >> vshift;\n",
      "      v.ui += (v.ui & 0x3fff) != 0x1000 || (significand & 0x7ff) ? 0x1000 : 0;\n",
      "    } else if (v.si <= maxN) {\n",
      "      // Handle norms\n",
      "      v.ui += (v.ui & 0x3fff) != 0x1000 ? 0x1000 : 0;\n",
      "      v.ui -= expAdjust << fp32FractionBits;\n",
      "    } else if (v.si <= infN) {\n",
      "      v.si = infN;\n",
      "    } else if (v.si < nanN) {\n",
      "      v.si = nanN;\n",
      "    }\n",
      "\n",
      "    v.ui >>= shift;\n",
      "    return sign | (v.ui & 0x7fff);\n",
      "  }\n",
      "\n",
      "  TVM_XINLINE float half2float(const uint16_t& value) const {\n",
      "    Bits v;\n",
      "    v.ui = value;\n",
      "    int32_t sign = v.si & signC;\n",
      "    v.si ^= sign;\n",
      "    sign <<= shiftSign;\n",
      "    v.si ^= ((v.si + minD) ^ v.si) & -(v.si > subC);\n",
      "    v.si ^= ((v.si + maxD) ^ v.si) & -(v.si > maxC);\n",
      "    Bits s;\n",
      "    s.si = mulC;\n",
      "    s.f *= v.si;\n",
      "    int32_t mask = -(norC > v.si);\n",
      "    v.si <<= shift;\n",
      "    v.si ^= (s.si ^ v.si) & mask;\n",
      "    v.si |= sign;\n",
      "    return v.f;\n",
      "  }\n",
      "\n",
      "  TVM_XINLINE float half2float(\n",
      "    const volatile uint16_t& value) const volatile {\n",
      "    Bits v;\n",
      "    v.ui = value;\n",
      "    int32_t sign = v.si & signC;\n",
      "    v.si ^= sign;\n",
      "    sign <<= shiftSign;\n",
      "    v.si ^= ((v.si + minD) ^ v.si) & -(v.si > subC);\n",
      "    v.si ^= ((v.si + maxD) ^ v.si) & -(v.si > maxC);\n",
      "    Bits s;\n",
      "    s.si = mulC;\n",
      "    s.f *= v.si;\n",
      "    int32_t mask = -(norC > v.si);\n",
      "    v.si <<= shift;\n",
      "    v.si ^= (s.si ^ v.si) & mask;\n",
      "    v.si |= sign;\n",
      "    return v.f;\n",
      "  }\n",
      "\n",
      "  template<typename T>\n",
      "  TVM_XINLINE void constructor(const T& value) {\n",
      "    half_ = float2half(float(value));\n",
      "  }\n",
      "};\n",
      "\n",
      "TVM_HALF_OPERATOR(half, +)\n",
      "TVM_HALF_OPERATOR(half, -)\n",
      "TVM_HALF_OPERATOR(half, *)\n",
      "TVM_HALF_OPERATOR(half, /)\n",
      "TVM_HALF_OPERATOR(bool, >)\n",
      "TVM_HALF_OPERATOR(bool, <)\n",
      "TVM_HALF_OPERATOR(bool, >=)\n",
      "TVM_HALF_OPERATOR(bool, <=)\n",
      "\n",
      "TVM_XINLINE half __float2half_rn(const float a) {\n",
      "  return half(a);\n",
      "}\n",
      "#else\n",
      "#include <cuda_fp16.h>\n",
      "__device__ half max(half a, half b)\n",
      "{\n",
      "  return __hgt(__half(a), __half(b)) ? a : b;\n",
      "}\n",
      "__device__ half min(half a, half b)\n",
      "{\n",
      "  return __hlt(__half(a), __half(b)) ? a : b;\n",
      "}\n",
      "#endif\n",
      "\n",
      "\n",
      "// Pack two half values.\n",
      "static inline __device__ __host__ unsigned\n",
      "__pack_half2(const half x, const half y) {\n",
      "  unsigned v0 = *((unsigned short *)&x);\n",
      "  unsigned v1 = *((unsigned short *)&y);\n",
      "  return (v1 << 16) | v0;\n",
      "}\n",
      "\n",
      "#define CUDA_UNSUPPORTED_HALF_MATH_BINARY(HALF_MATH_NAME, FP32_MATH_NAME) \\\n",
      "static inline __device__ __host__ half HALF_MATH_NAME(half x, half y) {   \\\n",
      "  float tmp_x = __half2float(x);                                          \\\n",
      "  float tmp_y = __half2float(y);                                          \\\n",
      "  float result = FP32_MATH_NAME(tmp_x, tmp_y);                            \\\n",
      "  return __float2half(result);                                            \\\n",
      "}\n",
      "\n",
      "#define CUDA_UNSUPPORTED_HALF_MATH_UNARY(HALF_MATH_NAME, FP32_MATH_NAME) \\\n",
      "static inline __device__ __host__ half HALF_MATH_NAME(half x) {          \\\n",
      "  float tmp_x = __half2float(x);                                         \\\n",
      "  float result = FP32_MATH_NAME(tmp_x);                                  \\\n",
      "  return __float2half(result);                                           \\\n",
      "}\n",
      "\n",
      "// Some fp16 math functions are not supported in cuda_fp16.h,\n",
      "// so we define them here to make sure the generated CUDA code\n",
      "// is valid.\n",
      "#if defined(__CUDA_ARCH__)\n",
      "#if (__CUDA_ARCH__ >= 530)\n",
      "CUDA_UNSUPPORTED_HALF_MATH_BINARY(hpow, powf)\n",
      "CUDA_UNSUPPORTED_HALF_MATH_UNARY(htanh, tanhf)\n",
      "CUDA_UNSUPPORTED_HALF_MATH_UNARY(htan, tanf)\n",
      "CUDA_UNSUPPORTED_HALF_MATH_UNARY(hatan, atanf)\n",
      "CUDA_UNSUPPORTED_HALF_MATH_UNARY(herf, erf)\n",
      "#else\n",
      "CUDA_UNSUPPORTED_HALF_MATH_UNARY(hexp, exp)\n",
      "#endif\n",
      "#endif\n",
      "\n",
      "#undef CUDA_UNSUPPORTED_HALF_MATH_BINARY\n",
      "#undef CUDA_UNSUPPORTED_HALF_MATH_UNARY\n",
      "\n",
      "struct __align__(8) half4 {\n",
      "  __half x, y, z, w;\n",
      "  __host__ __device__ half4() : x(__half(0)), y(__half(0)), z(__half(0)), w(__half(0)) {}\n",
      "  __host__ __device__ half4(__half x, __half y, __half z, __half w) : x(x), y(y), z(z), w(w) {}\n",
      "\n",
      "};\n",
      "__host__ __device__ half4 make_half4(__half x, __half y, __half z, __half w) {\n",
      "    return half4(x, y, z, w);\n",
      "}\n",
      "\n",
      "#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700)\n",
      "#define __shfl_sync(mask, var, lane, width) \\\n",
      "        __shfl((var), (lane), (width))\n",
      "\n",
      "#define __shfl_down_sync(mask, var, offset, width) \\\n",
      "        __shfl_down((var), (offset), (width))\n",
      "\n",
      "#define __shfl_up_sync(mask, var, offset, width) \\\n",
      "        __shfl_up((var), (offset), (width))\n",
      "#endif\n",
      "\n",
      "#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 610)\n",
      "#include <sm_61_intrinsics.h>\n",
      "\n",
      "\n",
      "#if defined(__CUDACC_RTC__)\n",
      "#define __SM_61_INTRINSICS_DECL__ __device__\n",
      "#else /* !__CUDACC_RTC__ */\n",
      "#define __SM_61_INTRINSICS_DECL__ static __device__ __inline__\n",
      "#endif /* __CUDACC_RTC__ */\n",
      "\n",
      "#ifndef __CUDA_ARCH__\n",
      "#define __DEF_IF_HOST { }\n",
      "#else  /* !__CUDA_ARCH__ */\n",
      "#define __DEF_IF_HOST ;\n",
      "#endif /* __CUDA_ARCH__ */\n",
      "\n",
      "__SM_61_INTRINSICS_DECL__ int __dp4a(unsigned int srcA, int srcB, int c) __DEF_IF_HOST\n",
      "__SM_61_INTRINSICS_DECL__ int __dp4a(int srcA, unsigned int srcB, int c) __DEF_IF_HOST\n",
      "\n",
      "#undef __DEF_IF_HOST\n",
      "\n",
      "#if !defined(__CUDACC_RTC__) && defined(__CUDA_ARCH__)\n",
      "__SM_61_INTRINSICS_DECL__ int __dp4a(unsigned int srcA, int srcB, int c) {\n",
      "    int ret;\n",
      "    asm volatile (\"dp4a.u32.s32 %0, %1, %2, %3;\" : \"=r\"(ret) : \"r\"(srcA), \"r\"(srcB), \"r\"(c));\n",
      "    return ret;\n",
      "}\n",
      "\n",
      "__SM_61_INTRINSICS_DECL__ int __dp4a(int srcA, unsigned int srcB, int c) {\n",
      "    int ret;\n",
      "    asm volatile (\"dp4a.s32.u32 %0, %1, %2, %3;\" : \"=r\"(ret) : \"r\"(srcA), \"r\"(srcB), \"r\"(c));\n",
      "    return ret;\n",
      "}\n",
      "#endif /* !__CUDACC_RTC__ && defined(__CUDA_ARCH__) */\n",
      "\n",
      "#undef __SM_61_INTRINSICS_DECL__\n",
      "\n",
      "#endif\n",
      "\n",
      "#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || \\\n",
      "     (__CUDACC_VER_MAJOR__ > 11))\n",
      "#define TVM_ENABLE_L2_PREFETCH 1\n",
      "#else\n",
      "#define TVM_ENABLE_L2_PREFETCH 0\n",
      "#endif\n",
      "\n",
      "#ifdef _WIN32\n",
      "  using uint = unsigned int;\n",
      "  using uchar = unsigned char;\n",
      "  using ushort = unsigned short;\n",
      "  using int64_t = long long;\n",
      "  using uint64_t = unsigned long long;\n",
      "#else\n",
      "  #define uint unsigned int\n",
      "  #define uchar unsigned char\n",
      "  #define ushort unsigned short\n",
      "  #define int64_t long long\n",
      "  #define uint64_t unsigned long long\n",
      "#endif\n",
      "\n",
      "#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 800) \n",
      "#define TVM_ENBALE_EFFICIENT_SMEM_PTR_CAST 1\n",
      "#else\n",
      "#define TVM_ENBALE_EFFICIENT_SMEM_PTR_CAST 0\n",
      "#endif\n",
      "extern \"C\" __global__ void __launch_bounds__(128) matmul_m1n16384k16384_f16xi4_simt_kernel(half* __restrict__ A, signed char* __restrict__ B, half* __restrict__ C);\n",
      "extern \"C\" __global__ void __launch_bounds__(128) matmul_m1n16384k16384_f16xi4_simt_kernel(half* __restrict__ A, signed char* __restrict__ B, half* __restrict__ C) {\n",
      "  half in_thread_C_local[1];\n",
      "  signed char B_local[4];\n",
      "  half B_decode_local[8];\n",
      "  half A_local[8];\n",
      "  __shared__ half red_result[1];\n",
      "  in_thread_C_local[0] = __float2half_rn(0.000000e+00f);\n",
      "  for (int ax1_0 = 0; ax1_0 < 16; ++ax1_0) {\n",
      "    *(int*)(B_local + 0) = *(int*)(B + (((((int)blockIdx.x) * 8192) + (ax1_0 * 512)) + (((int)threadIdx.x) * 4)));\n",
      "    for (int ax1 = 0; ax1 < 8; ++ax1) {\n",
      "      B_decode_local[ax1] = (((half)((((uint)B_local[(ax1 >> 1)]) >> (((uint)(ax1 & 1)) * (uint)4)) & (uint)15)) - __float2half_rn(8.000000e+00f));\n",
      "    }\n",
      "    *(uint4*)(A_local + 0) = *(uint4*)(A + ((ax1_0 * 1024) + (((int)threadIdx.x) * 8)));\n",
      "    for (int ax1_2_0 = 0; ax1_2_0 < 4; ++ax1_2_0) {\n",
      "      for (int ax1_2_1 = 0; ax1_2_1 < 2; ++ax1_2_1) {\n",
      "        in_thread_C_local[0] = (in_thread_C_local[0] + (A_local[((ax1_2_0 * 2) + ax1_2_1)] * B_decode_local[((ax1_2_0 * 2) + ax1_2_1)]));\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "  half red_buf0[1];\n",
      "  uint mask[1];\n",
      "  half t0[1];\n",
      "  half red_buf0_1[1];\n",
      "  uint mask_1[1];\n",
      "  half t0_1[1];\n",
      "  __shared__ half red_buf_staging[4];\n",
      "  red_buf0_1[0] = in_thread_C_local[0];\n",
      "  mask_1[0] = __activemask();\n",
      "  t0_1[0] = __shfl_down_sync(mask_1[0], red_buf0_1[0], 16, 32);\n",
      "  red_buf0_1[0] = (red_buf0_1[0] + t0_1[0]);\n",
      "  t0_1[0] = __shfl_down_sync(mask_1[0], red_buf0_1[0], 8, 32);\n",
      "  red_buf0_1[0] = (red_buf0_1[0] + t0_1[0]);\n",
      "  t0_1[0] = __shfl_down_sync(mask_1[0], red_buf0_1[0], 4, 32);\n",
      "  red_buf0_1[0] = (red_buf0_1[0] + t0_1[0]);\n",
      "  t0_1[0] = __shfl_down_sync(mask_1[0], red_buf0_1[0], 2, 32);\n",
      "  red_buf0_1[0] = (red_buf0_1[0] + t0_1[0]);\n",
      "  t0_1[0] = __shfl_down_sync(mask_1[0], red_buf0_1[0], 1, 32);\n",
      "  red_buf0_1[0] = (red_buf0_1[0] + t0_1[0]);\n",
      "  if ((((int)threadIdx.x) % 32) == 0) {\n",
      "    red_buf_staging[(((int)threadIdx.x) >> 5)] = red_buf0_1[0];\n",
      "  }\n",
      "  __syncthreads();\n",
      "  if (((int)threadIdx.x) < 4) {\n",
      "    red_buf0[0] = red_buf_staging[((int)threadIdx.x)];\n",
      "  }\n",
      "  mask[0] = __activemask();\n",
      "  t0[0] = __shfl_down_sync(mask[0], red_buf0[0], 2, 32);\n",
      "  red_buf0[0] = (red_buf0[0] + t0[0]);\n",
      "  t0[0] = __shfl_down_sync(mask[0], red_buf0[0], 1, 32);\n",
      "  red_buf0[0] = (red_buf0[0] + t0[0]);\n",
      "  if (((int)threadIdx.x) == 0) {\n",
      "    ((volatile half*)red_result)[0] = red_buf0[0];\n",
      "  }\n",
      "  __syncthreads();\n",
      "  if (((int)threadIdx.x) == 0) {\n",
      "    C[((int)blockIdx.x)] = (half)(((volatile half*)red_result)[0]);\n",
      "  }\n",
      "}\n",
      "\n",
      "\n",
      "extern \"C\" void init() {\n",
      "    \n",
      "}\n",
      "\n",
      "extern \"C\" void call(half* __restrict__ A, int8_t* __restrict__ B, half* __restrict__ C, cudaStream_t stream=cudaStreamDefault) {\n",
      "matmul_m1n16384k16384_f16xi4_simt_kernel<<<dim3(16384, 1, 1), dim3(128, 1, 1), 0, stream>>>(A, B, C);\n",
      "}\n",
      "\n",
      "latency=0.14325749999999998\n"
     ]
    }
   ],
   "source": [
    "matmul_config = bitblas.MatmulConfig(\n",
    "    M=1,  # M dimension\n",
    "    N=16384,  # N dimension\n",
    "    K=16384,  # K dimension\n",
    "    A_dtype=\"float16\",  # activation A dtype\n",
    "    W_dtype=\"int4\",  # weight W dtype\n",
    "    fast_decoding=True, # Disable Fast Decoding\n",
    ")\n",
    "\n",
    "\n",
    "matmul = bitblas.Matmul(config=matmul_config, enable_tuning=False)\n",
    "\n",
    "matmul.hardware_aware_finetune(topk=20, parallel_build=True)\n",
    "\n",
    "print(matmul.get_source())\n",
    "latency = matmul.profile_latency()\n",
    "print(f\"{latency=}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "feeeecb3-339d-42cc-9069-ee8645efc7d7",
   "metadata": {},
   "source": [
    "## Performance of Fast Dequantization on A100-80G\n",
    "\n",
    "![image.png](./img/FastDequantization_EXP.png)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
